#include "kernel_operator.h"
#include <stdio.h>

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

template<typename T> class KernelConcatd {
public:
    __aicore__ inline KernelConcatd() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, const ConcatdCustomTilingData &tiling_data)
    {
        this->tiling=tiling_data;
        if(AscendC::GetBlockIdx()<tiling.big_core_num){
        this->blockLength = tiling.big_tile_length;
        this->tileNum = tiling.big_tile_times;
        this->tailNum = tiling.big_tail_num;
        AscendC::PRINTF("blockLength string %d\n", this->blockLength);
        AscendC::PRINTF("tileNum string %d\n", this->tileNum );
        AscendC::PRINTF("tailNum string %d\n", this->tailNum);
        }else{
        this->blockLength = tiling.small_tile_length;
        this->tileNum = tiling.small_tile_times;
        this->tailNum = tiling.small_tail_num;
        AscendC::PRINTF("blockLength string %d\n", this->blockLength);
        AscendC::PRINTF("tileNum string %d\n", this->tileNum );
        AscendC::PRINTF("tailNum string %d\n", this->tailNum);
        }

        if(AscendC::GetBlockIdx()<tiling.big_core_num){
        xGm.SetGlobalBuffer((__gm__ T *)x + this->blockLength * AscendC::GetBlockIdx()*tiling.x2, this->blockLength*tiling.x2);
        yGm.SetGlobalBuffer((__gm__ T *)y + this->blockLength * AscendC::GetBlockIdx()*tiling.y2, this->blockLength*tiling.y2);
        zGm.SetGlobalBuffer((__gm__ T *)z + this->blockLength * AscendC::GetBlockIdx()*(tiling.x2+tiling.y2), this->blockLength*(tiling.x2+tiling.y2));
        }else{
        xGm.SetGlobalBuffer((__gm__ T *)x + (this->blockLength * AscendC::GetBlockIdx()+tiling.big_core_num)*tiling.x2, this->blockLength*tiling.x2);
        yGm.SetGlobalBuffer((__gm__ T *)y + (this->blockLength * AscendC::GetBlockIdx()+tiling.big_core_num)*tiling.y2, this->blockLength*tiling.y2);
        zGm.SetGlobalBuffer((__gm__ T *)z + (this->blockLength * AscendC::GetBlockIdx()+tiling.big_core_num)*(tiling.x2+tiling.y2), this->blockLength*(tiling.x2+tiling.y2));
        }

        AscendC::PRINTF("tailNum string %d\n", tiling.x2);
        AscendC::PRINTF("tailNum string %d\n", tiling.y2);
        AscendC::PRINTF("core_tile_x1 string %d\n", tiling.core_tile_x1);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, tiling.core_tile_x1* tiling.x2 * sizeof(T));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, tiling.core_tile_x1* tiling.y2 * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum;
        this->processNum = tiling.core_tile_x1;
        AscendC::PRINTF("fmt process %d\n", loopCount);
        for (int32_t i = 0; i < loopCount-1; i++) 
        {
            for(int32_t j = 0; j<this->processNum ;j++){
                CopyIn(i,j);
                CopyOut(i,j);
            }
        }
        this->processNum = this->tailNum;
        for(int32_t j =0;j<this->processNum ; j++){
            CopyIn(loopCount-1,j);
            CopyOut(loopCount-1,j);
        }
    } 

private:
    __aicore__ inline void CopyIn(int32_t progress,int32_t row)//
    {

        AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        AscendC::LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();

        AscendC::PRINTF("fmt process %d\n", progress);
        AscendC::PRINTF("fmt row %d\n", row);

        AscendC::DataCopyExtParams copyParams{1,  static_cast<uint32_t>(tiling.x2*sizeof(T)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
        AscendC::DataCopyPad(xLocal, xGm[progress *tiling.core_tile_x1 * tiling.x2 + tiling.x2 * row], copyParams,padParams); 

        AscendC::DataCopyExtParams copyParams0{1,  static_cast<uint32_t>(tiling.y2*sizeof(T)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<T> padParams0{true, 0, 0, 0};
        AscendC::DataCopyPad(yLocal, yGm[progress * tiling.core_tile_x1 * tiling.y2 + tiling.y2 * row], copyParams0,padParams0); 

        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    
    __aicore__ inline void CopyOut(int32_t progress,int32_t row)
    {
        AscendC::LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        AscendC::LocalTensor<T> yLocal = inQueueY.DeQue<T>();

        AscendC::PRINTF("fmt process %d\n", progress);
        AscendC::PRINTF("fmt row %d\n", row);

        AscendC::DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(T)*tiling.x2) , 0, static_cast<uint32_t>((tiling.y2)*sizeof(T)), 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPad(zGm[progress * (tiling.x2+tiling.y2)*tiling.core_tile_x1+(tiling.x2+tiling.y2)*row], xLocal, copyParams1); 

        AscendC::DataCopyExtParams copyParams2{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(T)*tiling.y2), 0, static_cast<uint32_t>((tiling.x2)*sizeof(T)), 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPad(zGm[(progress * (tiling.x2+tiling.y2)*tiling.core_tile_x1)+(tiling.x2+tiling.y2)*row + tiling.x2], yLocal, copyParams2); 

        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
    }



private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<T> xGm;
    AscendC::GlobalTensor<T> yGm;
    AscendC::GlobalTensor<T> zGm;
    AscendC::DataCopyPadExtParams<T> padParams;
    AscendC::DataCopyPadExtParams<T> padParams0;
    AscendC::DataCopyExtParams copyParams;
    AscendC::DataCopyExtParams copyParams0;
    AscendC::DataCopyExtParams copyParams1;
    AscendC::DataCopyExtParams copyParams2;
    AscendC::DataCopyExtParams copyParams3;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tailNum;
    uint32_t processNum;
    uint32_t tileLength;
    ConcatdCustomTilingData tiling;
};

template<typename T>class KernelConcatd0 {
public:
    __aicore__ inline KernelConcatd0() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, const ConcatdCustomTilingData &tiling_data)
    {
        this->tiling=tiling_data;
        if(AscendC::GetBlockIdx()<tiling.sbig_core_num){
        this->blockLength = tiling.sbig_tile_length;
        this->tileNum = tiling.sbig_tile_times;
        this->tailNum = tiling.sbig_tail_num;
        AscendC::PRINTF("blockLength string %d\n", this->blockLength);
        AscendC::PRINTF("tileNum string %d\n", this->tileNum );
        AscendC::PRINTF("tailNum string %d\n", this->tailNum);
        }else{
        this->blockLength = tiling.ssmall_tile_length;
        this->tileNum = tiling.ssmall_tile_times;
        this->tailNum = tiling.ssmall_tail_num;
        AscendC::PRINTF("blockLength string %d\n", this->blockLength);
        AscendC::PRINTF("tileNum string %d\n", this->tileNum );
        AscendC::PRINTF("tailNum string %d\n", this->tailNum);
        }

        if(AscendC::GetBlockIdx()<tiling.sbig_core_num){
            if(tiling.x1>this->blockLength * AscendC::GetBlockIdx()+this->blockLength-1){
                xGm.SetGlobalBuffer((__gm__ T *)x + this->blockLength * AscendC::GetBlockIdx()*tiling.x2, this->blockLength*tiling.x2);
                yGm.SetGlobalBuffer((__gm__ T *)y + this->blockLength * AscendC::GetBlockIdx()*tiling.y2, this->blockLength*tiling.y2);
                zGm.SetGlobalBuffer((__gm__ T *)z + this->blockLength * AscendC::GetBlockIdx()*tiling.x2, this->blockLength*tiling.x2);
            }else if(tiling.x1>this->blockLength * AscendC::GetBlockIdx()&&tiling.x1<this->blockLength * AscendC::GetBlockIdx()+this->blockLength){
                xGm.SetGlobalBuffer((__gm__ T *)x + this->blockLength * AscendC::GetBlockIdx()*tiling.x2, (tiling.x1-(this->blockLength * AscendC::GetBlockIdx()))*tiling.x2);
                yGm.SetGlobalBuffer((__gm__ T *)y + this->blockLength * AscendC::GetBlockIdx()*tiling.y2*0, (this->blockLength-(tiling.x1-(this->blockLength * AscendC::GetBlockIdx())))*tiling.y2);
                zGm.SetGlobalBuffer((__gm__ T *)z + this->blockLength * AscendC::GetBlockIdx()*tiling.x2, this->blockLength*tiling.x2);
            }else{
                xGm.SetGlobalBuffer((__gm__ T *)x + this->blockLength * AscendC::GetBlockIdx()*tiling.x2, this->blockLength*tiling.x2);
                yGm.SetGlobalBuffer((__gm__ T *)y + (this->blockLength * AscendC::GetBlockIdx()-tiling.x1) * tiling.y2,this->blockLength*tiling.y2);
                zGm.SetGlobalBuffer((__gm__ T *)z + this->blockLength * AscendC::GetBlockIdx()*tiling.x2, this->blockLength*tiling.x2);
            }
        }else{
            if(tiling.x1>this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num+this->blockLength-1){
                xGm.SetGlobalBuffer((__gm__ T *)x + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.x2, this->blockLength*tiling.x2);
                yGm.SetGlobalBuffer((__gm__ T *)y + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.y2, this->blockLength*tiling.y2);
                zGm.SetGlobalBuffer((__gm__ T *)z + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.x2, this->blockLength*tiling.x2);
            }else if(tiling.x1>this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num && tiling.x1<this->blockLength * AscendC::GetBlockIdx()+tiling.big_core_num+this->blockLength){
                xGm.SetGlobalBuffer((__gm__ T *)x + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.x2, (tiling.x1-(this->blockLength * AscendC::GetBlockIdx()+tiling.big_core_num))*tiling.x2);
                yGm.SetGlobalBuffer((__gm__ T *)y + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.y2*0, (this->blockLength-(tiling.x1-(this->blockLength * AscendC::GetBlockIdx()+tiling.big_core_num)))*tiling.y2);
                zGm.SetGlobalBuffer((__gm__ T *)z + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.x2, this->blockLength*tiling.x2);
            }else{
                xGm.SetGlobalBuffer((__gm__ T *)x + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.x2, this->blockLength*tiling.x2);
                yGm.SetGlobalBuffer((__gm__ T *)y + ((this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)-tiling.x1) * tiling.y2,this->blockLength*tiling.y2);
                zGm.SetGlobalBuffer((__gm__ T *)z + (this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num)*tiling.x2, this->blockLength*tiling.x2);
            }
        }

        AscendC::PRINTF("x1 string %d\n", tiling.x1);
        AscendC::PRINTF("y1 string %d\n", tiling.y1);
        AscendC::PRINTF("core_tile_s1 string %d\n", tiling.core_tile_s1);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, tiling.core_tile_s1* tiling.x2 * sizeof(T));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, tiling.core_tile_s1* tiling.y2 * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum;
        this->processNum = tiling.core_tile_s1;
        AscendC::PRINTF("fmt loopCount %d\n", loopCount);
        for (int32_t i = 0; i < loopCount-1; i++) 
        {
            for(int32_t j = 0; j<this->processNum ;j++){
                CopyIn(i,j);
                CopyOut(i,j);
            }
        }
        this->processNum = this->tailNum;
        for(int32_t j =0;j<this->processNum ; j++){
            CopyIn(loopCount-1,j);
            CopyOut(loopCount-1,j);
        }
    } 

private:
    __aicore__ inline void CopyIn(int32_t progress,int32_t row)//
    {

        AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        AscendC::LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();

        AscendC::PRINTF("fmt process %d\n", progress);
        AscendC::PRINTF("fmt row %d\n", row);
        int rowsss = 0;
        //判断是否搬完了x，就是要得到行数
        if(AscendC::GetBlockIdx()<tiling.sbig_core_num){
            rowsss = this->blockLength * AscendC::GetBlockIdx();
        }else{
            rowsss = this->blockLength * AscendC::GetBlockIdx() + tiling.sbig_core_num;
        }
        AscendC::PRINTF("hangshu row %d\n", rowsss + progress*tiling.core_tile_s1+row);
        if(rowsss + progress*tiling.core_tile_s1+row < tiling.x1){//如果行数小于x1，则随便搬x
            AscendC::DataCopyExtParams copyParams{1,  static_cast<uint32_t>(tiling.x2*sizeof(T)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
            AscendC::DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
            AscendC::DataCopyPad(xLocal, xGm[progress * tiling.core_tile_s1 * tiling.x2 + tiling.x2 * row], copyParams,padParams); 
        }else{//搬完x之后搬y
            //进入这个第一个数必为x1，判断x1是否还是在这个核心里面运行
            if(progress*tiling.core_tile_s1+row < this->blockLength  && rowsss<tiling.x1 && rowsss+this->blockLength>tiling.x1){//如果当前行数小于
            AscendC::DataCopyExtParams copyParams0{1,  static_cast<uint32_t>(tiling.y2*sizeof(T)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
            AscendC::DataCopyPadExtParams<T> padParams0{true, 0, 0, 0};
            AscendC::DataCopyPad(yLocal, yGm[(progress * tiling.core_tile_s1 + row - (tiling.x1-rowsss)) *tiling.y2  ], copyParams0,padParams0); 
            }else{
            //如果在对于y的搬法就得发生改变
            AscendC::DataCopyExtParams copyParams0{1,  static_cast<uint32_t>(tiling.y2*sizeof(T)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
            AscendC::DataCopyPadExtParams<T> padParams0{true, 0, 0, 0};
            AscendC::DataCopyPad(yLocal, yGm[progress * tiling.core_tile_s1 * tiling.y2 + tiling.y2 * row], copyParams0,padParams0); 
            }
        }

        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    
    __aicore__ inline void CopyOut(int32_t progress,int32_t row)
    {
        AscendC::LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        AscendC::LocalTensor<T> yLocal = inQueueY.DeQue<T>();

        AscendC::PRINTF("fmt process %d\n", progress);
        AscendC::PRINTF("fmt row %d\n", row);
        int rowsss = 0;
        if(AscendC::GetBlockIdx()<tiling.sbig_core_num){
            rowsss = this->blockLength * AscendC::GetBlockIdx();
        }else{
            rowsss = this->blockLength * AscendC::GetBlockIdx()+tiling.sbig_core_num;
        }
        //AscendC::PRINTF("hangshu row %d\n", rowsss + progress*tiling.core_tile_s1+row);
        if(rowsss + progress * tiling.core_tile_s1 + row < tiling.x1){//如果行数小于x1，则随便搬x
        AscendC::DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(T)*tiling.x2) , 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPad(zGm[progress * tiling.x2 * tiling.core_tile_s1 + tiling.x2*row], xLocal, copyParams1); 
        }else{//搬完搬y
        AscendC::DataCopyExtParams copyParams2{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(T)*tiling.y2), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPad(zGm[progress * tiling.y2 * tiling.core_tile_s1 + tiling.y2*row ], yLocal, copyParams2); 
        }
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
    }



private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<T> xGm;
    AscendC::GlobalTensor<T> yGm;
    AscendC::GlobalTensor<T> zGm;
    AscendC::DataCopyPadExtParams<T> padParams;
    AscendC::DataCopyPadExtParams<T> padParams0;
    AscendC::DataCopyExtParams copyParams;
    AscendC::DataCopyExtParams copyParams0;
    AscendC::DataCopyExtParams copyParams1;
    AscendC::DataCopyExtParams copyParams2;
    AscendC::DataCopyExtParams copyParams3;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tailNum;
    uint32_t processNum;
    uint32_t tileLength;
    ConcatdCustomTilingData tiling;
};

extern "C" __global__ __aicore__ void concatd_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace,  GM_ADDR tiling)
{

    GET_TILING_DATA(tiling_data, tiling);
    int32_t d = tiling_data.d;
    if(d==1){
        KernelConcatd<DTYPE_X> op;
        op.Init(x, y, z, tiling_data);
        op.Process();
    }
    if(d==0){
        KernelConcatd0<DTYPE_X> op;
        op.Init(x, y, z, tiling_data);
        op.Process();
    }
}